"""
1. 本题进行视频图像的目标跟踪。我们知道将一幅图像从BGR转换到颜色空间HSV，更容易表示一个特定颜色。在我们的程序中，要提取的是一个蓝色的口罩（kouzhao）。
窗口输出，参考下面图示。代码按下列要求完成：（35分）
"""
import cv2 as cv
import numpy as np

# ①    读取视频文件（video_mask.avi）
path = 'data/video_mask.avi'
cap = cv.VideoCapture(path)
W = int(cap.get(cv.CAP_PROP_FRAME_WIDTH))
H = int(cap.get(cv.CAP_PROP_FRAME_HEIGHT))
print('video size:', W, H)

# ②    利用（*'XVID'）定义4字符的编码器fourcc
fourcc = cv.VideoWriter_fourcc(*'XVID')

# ③    视频保存文件名output.avi，设置适当的帧率和视频尺度
FPS = 25
INTERVAL = 1000 // FPS
wt = cv.VideoWriter('output.avi', fourcc, FPS, (400-150, 475-300))

# ④    定义口罩蓝色的范围
mask_blue = np.uint8((136, 112, 82))  # BGR
mask_blue = np.expand_dims(mask_blue, axis=(0, 1))
mask_blue_hsv = cv.cvtColor(mask_blue, cv.COLOR_BGR2HSV)
mask_blue_h = mask_blue_hsv[0, 0, 0]

# ⑧    根据蓝色区间确定掩膜范围mask
range_lower = np.uint8([mask_blue_h - 30, 75, 75])
range_upper = np.uint8([mask_blue_h + 30, 255, 255])
print('range_lower', range_lower)
print('range_upper', range_upper)

# Keys for exiting
QUIT_KEYS = set([27, ord('q'), ord('Q')])

while True:
    # ⑤    读取视频帧，返回ret, frame
    ret, frame = cap.read()
    # 播放完视频自动退出
    if not ret:
        print('Video over, existing...')
        break

    # ⑥    在读取视频帧期间，确定帧图像口罩所在的适当区域frame_roi = frame[300:475, 150:400]
    frame_roi = frame[300:475, 150:400]

    # ⑦    将帧图像口罩所在区域frame_roi转换到HSV空间
    frame_roi_hsv = cv.cvtColor(frame_roi, cv.COLOR_BGR2HSV)

    # filter ROI HSV image
    mask_noise = cv.inRange(frame_roi_hsv, range_lower, range_upper)

    # ⑨    利用形态学开运算消除噪点
    kernel = np.ones((3, 3), dtype=np.uint8)
    mask = cv.morphologyEx(mask_noise, cv.MORPH_OPEN, kernel, iterations=3)

    # ⑩    利用形态学“膨胀”消除“空洞”
    mask = cv.dilate(mask, kernel, iterations=5)

    # 11利用位运算（ and ）还原口罩颜色
    filtered = cv.bitwise_and(frame_roi, frame_roi, mask=mask)

    # 12寻找口罩轮廓，确定矩形框
    contours, hierarchy = cv.findContours(mask, cv.RETR_EXTERNAL, cv.CHAIN_APPROX_SIMPLE)
    area_max = 0
    contour_i = 0
    for i, c in enumerate(contours):
        area = cv.contourArea(c)
        if area > area_max:
            area_max = area
            contour_i = i
    contour = contours[contour_i]

    # 13用绿色画出矩形框，并用白色显示字体（kouzhao）
    rect = cv.boundingRect(contour)
    c, r, w, h = rect
    cv.rectangle(filtered, (c, r), (c+w, r+h), (0, 255, 0), 2)
    cv.putText(filtered, 'kouzhao', (c, r - 10), cv.FONT_HERSHEY_PLAIN, 1.5, color=(255, 255, 255))

    # 14保存视频文件，并能够播放
    wt.write(filtered)

    # 15打开4个视频窗口，分别显示"original video"、"noise video"、"denoise video"、"dst video"。参考下面图示
    cv.imshow('original video', frame_roi)
    cv.imshow('noise video', mask_noise)
    cv.imshow('denoise video', mask)
    cv.imshow('dst video', filtered)

    # pause and get key
    k = cv.waitKey(INTERVAL) & 0xFF
    if k in QUIT_KEYS:
        print('Exit as user\'s demand.')
        break

# 16播放完视频自动退出，并释放关闭所有窗口
cap.release()
wt.release()
cv.destroyAllWindows()
